package com.duoec.dw.lineage.service.impl;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLObject;
import com.alibaba.druid.sql.ast.SQLOrderBy;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.SQLAssignItem;
import com.alibaba.druid.sql.ast.statement.SQLSelectGroupByClause;
import com.alibaba.druid.sql.ast.statement.SQLSelectQuery;
import com.alibaba.druid.sql.ast.statement.SQLWithSubqueryClause;
import com.duoec.dw.lineage.core.utils.AnalyserUtils;
import com.duoec.dw.lineage.dto.MultiTableLineageDto;
import com.duoec.dw.lineage.dto.TableFieldDependencyDto;
import com.duoec.dw.lineage.dto.TableFieldLineageDto;
import com.duoec.dw.lineage.dto.TableLineageDto;
import com.duoec.dw.lineage.exceptions.LineageAnalyseException;
import com.duoec.dw.lineage.service.HiveAnalyser;
import com.duoec.dw.lineage.service.analyser.*;
import com.google.common.collect.Maps;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import javax.annotation.PostConstruct;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * @author xuwenzhen
 */
@Service
public class HiveAnalyserImpl implements HiveAnalyser {
    @Autowired
    private List<HiveQueryAnalyser> hiveQueryAnalyserList;

    @Autowired
    private List<HiveStatementAnalyser<?>> hiveStatementAnalyserList;

    @Autowired
    private List<HiveFromAnalyser<?>> hiveFromAnalyserList;

    @Autowired
    private List<HiveExprAnalyser<?>> hiveExprAnalyserList;

    @Autowired
    private List<HiveSelectFieldAnalyser<?>> hiveSelectFieldAnalyserList;

    @Autowired
    private List<HiveOrderByAnalyser> hiveOrderByAnalyserList;

    @Autowired
    private List<HiveGroupByAnalyser> hiveGroupByAnalyserList;

    @Autowired
    private List<HivePartitionAnalyser> hivePartitionAnalyserList;

    private Map<Class<? extends SQLSelectQuery>, HiveQueryAnalyser<? extends SQLSelectQuery>> hiveQueryAnalyserMap;
    private Map<Class<? extends SQLStatement>, HiveStatementAnalyser<? extends SQLStatement>> hiveStatementAnalyserMap;
    private Map<Class<? extends SQLObject>, HiveFromAnalyser<? extends SQLObject>> hiveFromAnalyserMap;
    private Map<Class<? extends SQLExpr>, HiveExprAnalyser<? extends SQLExpr>> hiveExprAnalyserMap;
    private Map<Class<? extends SQLExpr>, HiveSelectFieldAnalyser<? extends SQLExpr>> hiveSelectFieldAnalyserMap;
    private Map<Class<SQLOrderBy>, HiveOrderByAnalyser> hiveOrderByAnalyserMap;
    private Map<Class<SQLSelectGroupByClause>, HiveGroupByAnalyser> hiveGroupByAnalyserMap;
    private Map<Class<SQLAssignItem>, HivePartitionAnalyser> hivePartitionAnalyserMap;

    @PostConstruct
    public void init() {
        if (CollectionUtils.isEmpty(hiveQueryAnalyserList)) {
            throw new LineageAnalyseException("无HiveSelectAnalyser实现类！");
        }
        if (CollectionUtils.isEmpty(hiveStatementAnalyserList)) {
            throw new LineageAnalyseException("无HiveStatementAnalyser实现类！");
        }
        if (CollectionUtils.isEmpty(hiveFromAnalyserList)) {
            throw new LineageAnalyseException("无HiveFromAnalyser实现类！");
        }
        if (CollectionUtils.isEmpty(hiveExprAnalyserList)) {
            throw new LineageAnalyseException("无HiveWhereAnalyser实现类！");
        }
        if (CollectionUtils.isEmpty(hiveSelectFieldAnalyserList)) {
            throw new LineageAnalyseException("无HiveSelectFieldAnalyser实现类！");
        }
        if (CollectionUtils.isEmpty(hiveOrderByAnalyserList)) {
            throw new LineageAnalyseException("无HiveOrderByAnalyserList实现类！");
        }
        if (CollectionUtils.isEmpty(hiveGroupByAnalyserList)) {
            throw new LineageAnalyseException("无HiveGroupByAnalyser实现类！");
        }
        if (CollectionUtils.isEmpty(hivePartitionAnalyserList)) {
            throw new LineageAnalyseException("无HivePartitionAnalyser实现类！");
        }
        hiveQueryAnalyserMap = hiveQueryAnalyserList
                .stream()
                .collect(Collectors.toMap(HiveQueryAnalyser::getSelectQueryType, analyser -> analyser));

        hiveStatementAnalyserMap = Maps.newHashMap();
        hiveStatementAnalyserList.forEach(analyser -> {
            Set<Class<? extends SQLStatement>> statementSet = analyser.parseSqlStatements();
            if (CollectionUtils.isEmpty(statementSet)) {
                return;
            }
            statementSet.forEach(statement -> hiveStatementAnalyserMap.put(statement, analyser));
        });

        hiveFromAnalyserMap = hiveFromAnalyserList
                .stream()
                .collect(Collectors.toMap(HiveFromAnalyser::getTableSourceType, analyser -> analyser));

        hiveExprAnalyserMap = hiveExprAnalyserList
                .stream()
                .collect(Collectors.toMap(HiveExprAnalyser::getSqlExprType, analyser -> analyser));

        hiveSelectFieldAnalyserMap = hiveSelectFieldAnalyserList
                .stream()
                .collect(Collectors.toMap(HiveSelectFieldAnalyser::getSelectFieldType, analyser -> analyser));

        hiveOrderByAnalyserMap = hiveOrderByAnalyserList
                .stream()
                .collect(Collectors.toMap(HiveOrderByAnalyser::getType, analyser -> analyser));

        hiveGroupByAnalyserMap = hiveGroupByAnalyserList
                .stream()
                .collect(Collectors.toMap(HiveGroupByAnalyser::getType, analyser -> analyser));

        hivePartitionAnalyserMap = hivePartitionAnalyserList
                .stream()
                .collect(Collectors.toMap(HivePartitionAnalyser::getType, analyser -> analyser));
    }

    /**
     * 分析 select 脚本
     *
     * @param sqlSelectQuery      查询脚本
     * @param withAsTableLineages with-as
     * @return 当前查询脚本的血缘关系数据
     */
    @Override
    public TableLineageDto analyseQuery(SQLSelectQuery sqlSelectQuery, MultiTableLineageDto withAsTableLineages) {
        HiveQueryAnalyser analyser = AnalyserUtils.getAnalyser(sqlSelectQuery, hiveQueryAnalyserMap);
        return analyser.analyse(sqlSelectQuery, withAsTableLineages);
    }

    /**
     * 分析statement
     *
     * @param statement SQLStatement
     * @return 当前操作的血缘关系
     */
    @Override
    public TableLineageDto analyseStatement(SQLStatement statement) {
        HiveStatementAnalyser statementAnalyser = AnalyserUtils.getAnalyser(statement, hiveStatementAnalyserMap);
        return statementAnalyser.analyse(statement);
    }

    /**
     * 分析from语句
     *
     * @param from                SQLTableSource
     * @param withAsTableLineages with-in
     * @return 当前操作的血缘关系
     */
    @Override
    public MultiTableLineageDto analyseFrom(SQLObject from, MultiTableLineageDto withAsTableLineages) {
        HiveFromAnalyser fromAnalyser = AnalyserUtils.getAnalyser(from, hiveFromAnalyserMap);
        return fromAnalyser.analyse(from, withAsTableLineages);
    }

    /**
     * 分析where语句
     *
     * @param fromLineages 当前语句from的血缘关系（多个）
     * @param where        where语句
     * @return 解析此表达式可抽取的字段依赖
     */
    @Override
    public LinkedHashSet<TableFieldDependencyDto> analyseWhere(MultiTableLineageDto fromLineages, SQLExpr where) {
        return analyseExpr(fromLineages, where);
    }

    /**
     * 分析表达式语句
     *
     * @param fromLineages 当前语句from的血缘关系（多个）
     * @param expr         表达式语句
     */
    @Override
    public LinkedHashSet<TableFieldDependencyDto> analyseExpr(MultiTableLineageDto fromLineages, SQLExpr expr) {
        if (expr == null) {
            return null;
        }
        HiveExprAnalyser exprAnalyser = AnalyserUtils.getAnalyser(expr, hiveExprAnalyserMap);
        return exprAnalyser.analyse(expr, fromLineages);
    }

    /**
     * 分析order by语句
     *
     * @param tableLineage      当前语句的血缘关系
     * @param fromTableLineages 当前语句from的血缘关系（多个）
     * @param orderBy           orderBy语句
     */
    @Override
    public void analyseOrderBy(TableLineageDto tableLineage, MultiTableLineageDto fromTableLineages, SQLOrderBy orderBy) {
        if (orderBy == null || CollectionUtils.isEmpty(orderBy.getItems())) {
            return;
        }
        HiveOrderByAnalyser analyser = AnalyserUtils.getAnalyser(orderBy, hiveOrderByAnalyserMap);
        analyser.analyse(orderBy, fromTableLineages, tableLineage);
    }

    /**
     * 分析group by语句
     *
     * @param tableLineage      当前语句的血缘关系
     * @param fromTableLineages 当前语句from的血缘关系（多个）
     * @param groupBy           group by语句
     */
    @Override
    public void analyseGroupBy(TableLineageDto tableLineage, MultiTableLineageDto fromTableLineages, SQLSelectGroupByClause groupBy) {
        if (groupBy == null || CollectionUtils.isEmpty(groupBy.getItems())) {
            return;
        }
        HiveGroupByAnalyser analyser = AnalyserUtils.getAnalyser(groupBy, hiveGroupByAnalyserMap);
        analyser.analyse(groupBy, fromTableLineages, tableLineage);
    }

    /**
     * partitions语句分析
     *
     * @param tableLineage 当前语句的血缘关系
     * @param partitions   partitions语句
     */
    @Override
    public void analysePartitions(TableLineageDto tableLineage, List<SQLAssignItem> partitions) {
        if (partitions == null || CollectionUtils.isEmpty(partitions)) {
            return;
        }
        HivePartitionAnalyser analyser = AnalyserUtils.getAnalyser(partitions.get(0), hivePartitionAnalyserMap);
        partitions.forEach(partition -> analyser.analyse(partition, tableLineage));
    }

    /**
     * 处理select字段表达式
     *
     * @param tableLineage      当前语句的血缘关系
     * @param fromTableLineages 当前语句from的血缘关系（多个）
     * @param selectExpr        select字段表达式
     * @return 当前解析出来的血缘关系字段(正常情况下返回一个 ， 如果是 * 时 ， 可以返回多个 ！)
     */
    @Override
    public LinkedHashSet<TableFieldLineageDto> analyseSelect(TableLineageDto tableLineage, MultiTableLineageDto fromTableLineages, SQLExpr selectExpr) {
        HiveSelectFieldAnalyser analyser = AnalyserUtils.getAnalyser(selectExpr, hiveSelectFieldAnalyserMap);
        return analyser.analyse(selectExpr, tableLineage, fromTableLineages);
    }

    public void setHiveQueryAnalyserList(List<HiveQueryAnalyser> hiveQueryAnalyserList) {
        this.hiveQueryAnalyserList = hiveQueryAnalyserList;
    }

    public void setHiveStatementAnalyserList(List<HiveStatementAnalyser<?>> hiveStatementAnalyserList) {
        this.hiveStatementAnalyserList = hiveStatementAnalyserList;
    }

    public void setHiveFromAnalyserList(List<HiveFromAnalyser<?>> hiveFromAnalyserList) {
        this.hiveFromAnalyserList = hiveFromAnalyserList;
    }

    public void setHiveExprAnalyserList(List<HiveExprAnalyser<?>> hiveExprAnalyserList) {
        this.hiveExprAnalyserList = hiveExprAnalyserList;
    }

    public void setHiveSelectFieldAnalyserList(List<HiveSelectFieldAnalyser<?>> hiveSelectFieldAnalyserList) {
        this.hiveSelectFieldAnalyserList = hiveSelectFieldAnalyserList;
    }

    public void setHiveOrderByAnalyserList(List<HiveOrderByAnalyser> hiveOrderByAnalyserList) {
        this.hiveOrderByAnalyserList = hiveOrderByAnalyserList;
    }

    public void setHiveGroupByAnalyserList(List<HiveGroupByAnalyser> hiveGroupByAnalyserList) {
        this.hiveGroupByAnalyserList = hiveGroupByAnalyserList;
    }

    public void setHivePartitionAnalyserList(List<HivePartitionAnalyser> hivePartitionAnalyserList) {
        this.hivePartitionAnalyserList = hivePartitionAnalyserList;
    }
}
